/* Copyright 2012 Tobias Marschall
 * 
 * This file is part of CLEVER.
 * 
 * CLEVER is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * CLEVER is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <iostream>
#include <fstream>
#include <vector>
#include <cassert>

#include <boost/program_options.hpp>
#include <boost/unordered_set.hpp>
#include <boost/math/distributions/normal.hpp>

#include "AlignmentPair.h"
#include "ReadAssignmentFileReader.h"
#include "VariationIndex.h"
#include "LengthAwareVariationIndex.h"
#include "VariationListParser.h"

using namespace std;
using namespace boost;
namespace po = boost::program_options;

long long skipped_duplicates = 0;
const double sqrt2 = 1.4142135623730951;

void usage(const char* name, const po::options_description& options_desc) {
	cerr << "Usage: " << name << " [options] <variants-file>" << endl;
	cerr << endl;
	cerr << "Reads alignment priors from stdin and extracts those that support" << endl;
	cerr << "at list one of the given variants. For all these reads, a SMART input" << endl;
	cerr << "file is created." << endl;
	cerr << endl;
	cerr << "NOTE: Input alignment priors must be grouped by read names, that is," << endl;
	cerr << "      different alignment pairs for the same read must be given in" << endl;
	cerr << "      consecutive lines." << endl;
	cerr << endl;
	cerr << options_desc << endl;
	exit(1);
}

typedef struct {
	int alignment_id;
	int variant_id;
	int insert_length;
	double alignment_prob;
	double initial_prob;
	// probability that this pairing is correct, given that the alignment is correct,
	// i.e. distribution over all variants the current alignment maps to
	double intra_alignment_prob;
} read_assignment_t;

// Verify that probabilities sum up to one for the given set of alignment pairs.
bool check_probability_sums(const vector<AlignmentPair>& pairs) {
	double psum1 = 0.0;
	double psum2 = 0.0;
	vector<AlignmentPair>::const_iterator it = pairs.begin();
	for (;it!=pairs.end(); ++it) {
		psum1 += it->getProbability();
		psum2 += it->getProbabilityInsertLength();
	}
	return (fabs(1.0-psum1)<1e-6) && (fabs(1.0-psum2)<1e-6);
}

/** Processes the top read in the input. */
void process_read(const vector<read_assignment_t>& assignments, const vector<AlignmentPair>& pairs, const vector<size_t>& assignment_counts, boost::math::normal& normal) {
	assert(pairs.size()>0);
	// Don't output anything for reads that do not map to at least one variation
	if (assignments.size() == 0) return;
	if (!check_probability_sums(pairs)) {
		cerr << "Error: probabilities do not sum to one for read \"" << pairs[0].getName() << "\"" << endl;
		exit(1);
	}
	int phred_sum = pairs[0].getPhredSum1() + pairs[0].getPhredSum2();
	cout << pairs[0].getName() << '\t' << phred_sum;
	// all the alignment indices that are associated with a variant.
	boost::unordered_set<int> used_indices;
	typedef vector<read_assignment_t>::const_iterator assignment_it;
	for (assignment_it it=assignments.begin(); it!=assignments.end(); ++it) {
		used_indices.insert(it->alignment_id);
	}
	// compute probabilities for "no-variant" component
	double no_variant_density = 0.0;
	double no_variant_posterior = 0.0;
	for (size_t i=0; i<pairs.size(); ++i) {
		if (used_indices.find(i) != used_indices.end()) continue;
		const AlignmentPair& ap = pairs[i];
		int insert_size = ap.getStart2() - ap.getEnd1() - 1;
		no_variant_density += ap.getProbability() * pdf(normal, insert_size);
		no_variant_posterior += ap.getProbabilityInsertLength();
	}
	cout << '\t' << no_variant_density << '\t' << no_variant_posterior;
	for (assignment_it it=assignments.begin(); it!=assignments.end(); ++it) {
		cout << '\t' << it->variant_id << ',' << it->alignment_id;
		cout << ',' << it->insert_length << ',' <<  it->alignment_prob;
		cout << ',' <<  it->initial_prob * it->intra_alignment_prob;
	}
	cout << endl;
}

int main(int argc, char* argv[]) {
	// PARAMETERS
	double insert_mean;
	double insert_stddev;
	double significance_level;

	po::options_description options_desc("Allowed options");
	options_desc.add_options()
		("insert_mean,i", po::value<double>(&insert_mean)->default_value(112.0), "Mean distance between two paired reads (excluding the reads).")
		("insert_stddev,d", po::value<double>(&insert_stddev)->default_value(11.0*1.4826), "Standard deviation of insert size (distance between paired reads) distribution, default: 11 * 1.4826 = median absolute deviation * correction factor)")
		("significance_level,p", po::value<double>(&significance_level)->default_value(0.05), "Level of evidence necessary NOT to associate read with a enclosed variation.")
	;

	if (isatty(fileno(stdin)) || (argc<2)) {
		usage(argv[0], options_desc);
	}
	string variants_filename(argv[argc-1]);
	argc -= 1;

	po::variables_map options;
	try {
		po::store(po::parse_command_line(argc, argv, options_desc), options);
		po::notify(options);
	} catch(std::exception& e) {
		cerr << "error: " << e.what() << "\n";
		return 1;
	}

	ifstream variants_stream(variants_filename.c_str());
	if (variants_stream.fail()) {
		cerr << "Error: could not open \"" << variants_filename << "\"." << endl;
		return 1;
	}
	auto_ptr<vector<Variation> > variations = VariationListParser::parse(variants_stream, true);
	cerr << "Read " << variations->size() << " variations." << endl;
 	LengthAwareVariationIndex* variation_index = new LengthAwareVariationIndex(*variations.get());

	boost::math::normal normal(insert_mean,insert_stddev);

	// TODO: double-check sqrt2 factor
 	boost::math::normal normal2(0,sqrt2*insert_stddev);
 	int allowable_insert_size_diff = (int)floor(-boost::math::quantile(normal2, significance_level/2.0));

	unsigned int n = 0;
	string last_read_name = "";
	// all the alignment pairs belonging to the currently processed read
	vector<AlignmentPair> alignment_pairs;
	// for each alignment pair in alignment_pairs, the number of variants it is associated to
	vector<size_t> assignment_counts;
	// pairs (variation_id, alignment_pair_nr)
	vector<read_assignment_t> assignments;
	string line;
	while (getline(cin, line)) {
		n += 1;
		if (n % 200000 == 0) {
			cerr << "Having processed " << n << " alignment pairs." << endl;
		}
		try {
			AlignmentPair ap(line);
			if (ap.getChrom1().compare(ap.getChrom2()) != 0) continue;
			if ((last_read_name.compare(ap.getName()) != 0) && (alignment_pairs.size() > 0)) {
				process_read(assignments, alignment_pairs, assignment_counts, normal);
				alignment_pairs.clear();
				assignment_counts.clear();
				assignments.clear();
			}
			alignment_pairs.push_back(ap);
			// determine the associated variations
			int insert_length_diff = ((int)ap.getStart2()) - ((int)ap.getEnd1()) - 1 - insert_mean;
			int min_length = insert_length_diff - allowable_insert_size_diff;
			int max_length = insert_length_diff + allowable_insert_size_diff;
			auto_ptr<vector<size_t> > variation_indices = variation_index->containedIn(ap.getChrom1(), ap.getEnd1()+1, ap.getStart2(), min_length, max_length);
			if (variation_indices.get() != 0) {
				double psum = 0.0;
				size_t first = assignments.size();
				for (vector<size_t>::const_iterator it = variation_indices->begin(); it != variation_indices->end(); ++it) {
					const Variation& v = variations->at(*it);
					// check whether length is compatible with considered variation
					assert(abs(insert_length_diff - v.getLengthDifference()) <= allowable_insert_size_diff);
					read_assignment_t assignment;
					assignment.alignment_id = ap.getPairNr();
					assignment.variant_id = *it;
					assignment.insert_length = ap.getStart2() - ap.getEnd1() - 1;
					assignment.alignment_prob = ap.getProbability();
					assignment.initial_prob = ap.getProbabilityInsertLength();
					assignment.intra_alignment_prob = pdf(normal,abs(insert_length_diff - v.getLengthDifference()));
					psum += assignment.intra_alignment_prob;
					// cerr << "Read " << ap.getName() << " (" << insert_length << ") is associated with variation " << *it << " (" << v.getLengthDifference() << ")" << endl;
					assignments.push_back(assignment);
				}
				// normalize intra alignment probabilities
				for (size_t i=first; i<assignments.size(); ++i) {
					assignments[i].intra_alignment_prob /= psum;
				}
				assignment_counts.push_back(variation_indices->size());
			} else {
				assignment_counts.push_back(0);
			}
			last_read_name = ap.getName();
		} catch (std::runtime_error&) {
			cerr << "Error parsing input, offending line: " << n << endl;
			return 1;
		}
	}
	if (alignment_pairs.size() > 0) {
		process_read(assignments, alignment_pairs, assignment_counts, normal);
	}

	delete variation_index;
	return 0;
}
